// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

struct TaskCompletedEvent {
  std::string taskUuid;
  exec::TaskState state;
  std::exception_ptr error;
  exec::TaskStats stats;

  std::string errorMessage() const {
    if (error) {
      try {
        std::rethrow_exception(error);
      } catch (const PolluxException& e) {
        return e.message();
      } catch (const std::exception& e) {
        return e.what();
      }
    } else {
      return "";
    }
  }
};

class TestTaskListener : public exec::TaskListener {
 public:
  void onTaskCompletion(
      const std::string& taskUuid,
      const std::string& taskId,
      exec::TaskState state,
      std::exception_ptr error,
      exec::TaskStats stats) override {
    events_.push_back({taskUuid, state, error, stats});
  }

  std::vector<TaskCompletedEvent>& events() {
    return events_;
  }

 private:
  std::vector<TaskCompletedEvent> events_;
};

class TaskListenerTest : public OperatorTestBase {};

TEST_F(TaskListenerTest, success) {
  auto data = make_row_vector({make_flat_vector<int32_t>({0, 1, 2, 3, 4})});

  auto plan = PlanBuilder().values({data}).planNode();

  // Register event listener to collect task completion events.
  auto listener = std::make_shared<TestTaskListener>();
  auto& events = listener->events();
  ASSERT_TRUE(exec::registerTaskListener(listener));

  assertQuery(plan, "VALUES (0), (1), (2), (3), (4)");
  ASSERT_EQ(1, events.size());
  ASSERT_EQ(nullptr, events.back().error);

  // Register the same listener again. This should have no effect as the
  // registration logic should detect a duplicate and not register it again.
  ASSERT_FALSE(exec::registerTaskListener(listener));

  // Clear the events, re-run the query and verify that a new event is received.
  events.clear();

  assertQuery(plan, "VALUES (0), (1), (2), (3), (4)");
  ASSERT_EQ(1, events.size());
  ASSERT_EQ(nullptr, events.back().error);

  // Clear the events, unregister the listener, re-run the query and verify that
  // no event is received.
  events.clear();
  ASSERT_TRUE(exec::unregisterTaskListener(listener));

  assertQuery(plan, "VALUES (0), (1), (2), (3), (4)");
  ASSERT_TRUE(events.empty());

  // Try to unregister the listener again.
  ASSERT_FALSE(exec::unregisterTaskListener(listener));
}

TEST_F(TaskListenerTest, error) {
  auto data = make_row_vector({make_flat_vector<int32_t>({0, 1, 2, 3, 4})});

  auto plan = PlanBuilder().values({data}).project({"10 / c0"}).planNode();

  CursorParameters params;
  params.planNode = plan;

  // Register event listener to collect task completion events.
  auto listener = std::make_shared<TestTaskListener>();
  auto& events = listener->events();
  ASSERT_TRUE(exec::registerTaskListener(listener));

  EXPECT_THROW(readCursor(params, [](auto) {}), PolluxException);

  ASSERT_EQ(1, events.size());
  ASSERT_EQ("division by zero", events.back().errorMessage());

  ASSERT_TRUE(exec::unregisterTaskListener(listener));
}

TEST_F(TaskListenerTest, multipleListeners) {
  auto listener1 = std::make_shared<TestTaskListener>();
  auto listener2 = std::make_shared<TestTaskListener>();
  ASSERT_TRUE(exec::registerTaskListener(listener1));
  ASSERT_TRUE(exec::registerTaskListener(listener2));

  ASSERT_TRUE(listener1->events().empty());
  ASSERT_TRUE(listener2->events().empty());

  auto data = make_row_vector({make_flat_vector<int32_t>({0, 1, 2, 3, 4})});
  auto plan = PlanBuilder().values({data}).planNode();
  assertQuery(plan, "VALUES (0), (1), (2), (3), (4)");

  ASSERT_EQ(1, listener1->events().size());
  ASSERT_EQ(1, listener2->events().size());

  {
    const auto& event1 = listener1->events().front();
    const auto& event2 = listener2->events().front();
    ASSERT_FALSE(event1.taskUuid.empty());
    ASSERT_EQ(event1.taskUuid, event2.taskUuid);
    ASSERT_EQ(event1.state, exec::TaskState::kFinished);
    ASSERT_EQ(event2.state, exec::TaskState::kFinished);
  }

  ASSERT_TRUE(exec::unregisterTaskListener(listener1));
  ASSERT_TRUE(exec::unregisterTaskListener(listener2));
}
